package inference;


import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.gps.utils.LatLonPoint;

import utils.datetime.DateTime;
import utils.kml.Kml;
import visual.KMLCirclePlacemark;
import visual.KMLBarChart;
import parser.TwitterParser;

import com.aliasi.cluster.LatentDirichletAllocation;

import com.aliasi.corpus.ObjectHandler;

import com.aliasi.stats.Statistics;
import com.aliasi.symbol.SymbolTable;
import com.aliasi.util.ObjectToCounterMap;
import com.aliasi.util.ObjectToDoubleMap;
import com.aliasi.util.ScoredObject;

import data.TopicSim;


public class LingPipeLdaHandler implements ObjectHandler<LatentDirichletAllocation.GibbsSample> {
	public SymbolTable mSymbolTable;
	public TwitterParser par;
	public int cellsize;
	public PrintWriter kml;
	public PrintWriter csv;
	
	//ordering types: NB: only of three
	boolean prob = true;
	boolean z = false;
	boolean p_z = false;
	
	//print variables
	boolean printDoc = false;
	boolean printWords = true;
	int maxWords = 5;


    public LingPipeLdaHandler(SymbolTable symbolTable, TwitterParser parser, int size, PrintWriter out, PrintWriter outcsv) {
    	mSymbolTable = symbolTable;
    	par = parser;
        cellsize = size;
        kml = out;
        csv = outcsv;
    }

    public void handle(LatentDirichletAllocation.GibbsSample sample) {
        System.out.println("numDocuments=" + sample.numDocuments());
        System.out.println("numTokens=" + sample.numTokens());
        System.out.println("numWords=" + sample.numWords());
        System.out.println("numTopics=" + sample.numTopics());
    }
    
	public void report(LatentDirichletAllocation.GibbsSample sample, int topic){
		Map<String,Double> labelKeyProb = new HashMap<String,Double>();
		Map<String,Integer> labelKeyHour = new HashMap<String,Integer>();
		findLabels(sample, topic, maxWords, labelKeyProb, labelKeyHour);
		reportTopicDocs(sample, topic, printDoc);
		reportTopicWords(sample, topic, printWords);
		printKml(sample, topic, kml, labelKeyProb, labelKeyHour);
		printCsv(sample, topic, labelKeyHour);
	}
    
	public List<TopicSim> similarity(LatentDirichletAllocation lda1, LatentDirichletAllocation lda2){
		//given 2 lda models, this function returns an ordered list of topics according to their similarity
		
		int numTopics = lda1.numTopics();
		List<TopicSim> pairs = new ArrayList<TopicSim>();
		List<TopicSim> listSim = new ArrayList<TopicSim>();
		
		for(int i=0; i<numTopics; i++){
			double [] pi = lda1.wordProbabilities(i);
			for(int j=0; j<numTopics;j++){
				double [] pj = lda2.wordProbabilities(j);
				double divergence = Statistics.symmetrizedKlDivergence(pi,pj);
				TopicSim ts = new TopicSim(i,j,divergence);
				pairs.add(ts);				
			}
			
		}
		
		Collections.sort(pairs, ScoredObject.comparator());
		boolean[] taken1 = new boolean[numTopics];
		boolean[] taken2 = new boolean[numTopics];
		
		for(TopicSim ts: pairs){
			if(!taken1[ts.mI] && !taken2[ts.mJ]){
				taken1[ts.mI] = taken2[ts.mJ] = true;
				listSim.add(ts);
			}
		}
		
		return listSim;
		
	}
	
	public void findLabels(LatentDirichletAllocation.GibbsSample sample, int topic, int topWords, Map<String,Double> labelKeyProb, Map<String,Integer> labelKeyHour){
		
		boolean [] hourTaken = new boolean[par.TIMEINTERVALS];
		double maxp = Double.MIN_VALUE;
		double minp = Double.MAX_VALUE;
		
		List<Integer> listWords = new ArrayList<Integer>();
		
		if(prob){
	        ObjectToCounterMap<Integer> counter = new ObjectToCounterMap<Integer>();
	        for (int word = 0; word < sample.numWords(); word++){
	        	counter.set(word, sample.topicWordCount(topic,word));
				maxp = Math.max(maxp, sample.topicWordProb(topic, word));
				minp = Math.min(minp, sample.topicWordProb(topic, word));
	        }
	        listWords = counter.keysOrderedByCountList();
	        
	        for(int rank=0;rank<topWords && rank<listWords.size();rank++){
	        	int wordId = listWords.get(rank);
	        	String word = mSymbolTable.idToSymbol(wordId);

	            String[] token = TwitterParser.translateWordArray(word);
	            if(!hourTaken[Integer.parseInt(token[0])]){
	                hourTaken[Integer.parseInt(token[0])] = true;
	                String key = token[1].trim()+"-"+token[2].trim();
	                double p = ((sample.topicWordProb(topic, wordId)-minp)/(maxp-minp))*((double)cellsize/(double)2);
	                labelKeyProb.put(key, p);
	                labelKeyHour.put(key, Integer.parseInt(token[0]));
	            }
	        		
	         }
		}
		
		if(z){
			int topicCount = sample.topicCount(topic);
    		ObjectToDoubleMap<Integer> counter = new ObjectToDoubleMap<Integer>();
    		for (int word = 0; word < sample.numWords(); word++){
    			int topicWordCount = sample.topicWordCount(topic,word);
    			int wordCount = sample.wordCount(word);
    			double z = binomialZ(topicWordCount, topicCount, wordCount, sample.numTokens());
    			counter.set(word,z);
    			maxp = Math.max(maxp, z);
    			minp = Math.min(minp, z);
    		}
    		
    		listWords = counter.keysOrderedByValueList();
    		
            for(int rank=0;rank<topWords && rank<listWords.size();rank++){
	        	int wordId = listWords.get(rank);
	            String word = mSymbolTable.idToSymbol(wordId);
	            String[] token = TwitterParser.translateWordArray(word);
	            if(!hourTaken[Integer.parseInt(token[0])]){
	                hourTaken[Integer.parseInt(token[0])] = true;
	                String key = token[1].trim()+"-"+token[2].trim();
		            int wordCount = sample.wordCount(wordId);
		            int topicWordCount = sample.topicWordCount(topic,wordId);
		            
		            double z = binomialZ(topicWordCount, topicCount, wordCount, sample.numTokens());
		            double p = ((z-minp)/(maxp-minp))*((double)cellsize/(double)2);
	                labelKeyProb.put(key, p);
	                labelKeyHour.put(key, Integer.parseInt(token[0]));
	            
	            }
            		
             }
		}
		
		
		//order word according to their z probability
		if(p_z){
			int topicCount = sample.topicCount(topic);
			
    		ObjectToDoubleMap<Integer> counter = new ObjectToDoubleMap<Integer>();
    		for (int word = 0; word < sample.numWords(); word++){
    			int topicWordCount = sample.topicWordCount(topic,word);
    			int wordCount = sample.wordCount(word);
    			double z = binomialZ(topicWordCount, topicCount, wordCount, sample.numTokens());
    			
    			double p = (double)sample.topicWordProb(topic, word);
    			double all = z*p;
    			counter.set(word,all);
    			maxp = Math.max(maxp, all);
    			minp = Math.min(minp, all);
    		}
    		
    		listWords = counter.keysOrderedByValueList();
    		    	        
    		for(int rank=0;rank<topWords && rank<listWords.size();rank++){
		    	
	        	int wordId = listWords.get(rank);
	            String word = mSymbolTable.idToSymbol(wordId);
	            String[] token = TwitterParser.translateWordArray(word);
	            if(!hourTaken[Integer.parseInt(token[0])]){
	                hourTaken[Integer.parseInt(token[0])] = true;
	                String key = token[1].trim()+"-"+token[2].trim();
		            int wordCount = sample.wordCount(wordId);
		            int topicWordCount = sample.topicWordCount(topic,wordId);
		            double topicWordProb = (double)sample.topicWordProb(topic,wordId);
		            double z = binomialZ(topicWordCount, topicCount, wordCount, sample.numTokens());
		            double all = z*topicWordProb;
		            double p = ((all-minp)/(maxp-minp))*((double)cellsize/(double)2);
		            labelKeyProb.put(key, p);
	                labelKeyHour.put(key, Integer.parseInt(token[0]));
	            }
	        }
		}
	}
    

    
    public void reportTopicDocs(LatentDirichletAllocation.GibbsSample sample, int topic, boolean reportDoc) {
    	if(reportDoc){
    		ObjectToCounterMap<Integer> counter = new ObjectToCounterMap<Integer>();
    		for (int doc = 0; doc < sample.numDocuments(); doc++) 
    			counter.set(doc, sample.documentTopicCount(doc,topic));
    		
    		List<Integer> topDocs = counter.keysOrderedByCountList();
            System.out.println();
            System.out.println("TOPIC: "+topic);
    		
            System.out.println("DOC      NAME                  COUNT    PROB");
            System.out.println("--------------------------------------------");
            for (int rank = 0; rank < topDocs.size(); rank++) {
                int doc = topDocs.get(rank);
                int docTopicCount = sample.documentTopicCount(doc,topic);
                double docTopicProb = sample.documentTopicProb(doc,topic);
                System.out.printf("%3d  %s   %7d   %4.3f\n", doc, par.int2filename.get(doc), docTopicCount, docTopicProb);
            }	
    	}
    }
    
    public void reportTopicWords(LatentDirichletAllocation.GibbsSample sample, int topic, boolean print) {
    	if(print){
    		//order word according to their probability
    		if(prob){
	    		ObjectToCounterMap<Integer> counter = new ObjectToCounterMap<Integer>();
	    		for (int word = 0; word < sample.numWords(); word++) 
	    			counter.set(word, sample.topicWordCount(topic, word));
	    		
	    		List<Integer> listWords = counter.keysOrderedByCountList();
	    		int topicCount = sample.topicCount(topic);
	    		
		        System.out.println("\nTOPIC " + topic  + "  (total count=" + topicCount + ")");
		        System.out.println("SYMBOL             WORD    COUNT   PROB     Z");
		        System.out.println("---------------------------------------------");
		            	        
		        for (int rank = 0; rank < listWords.size(); rank++) {
			    	
		        	int wordId = listWords.get(rank);
		            String word = mSymbolTable.idToSymbol(wordId);
		            int wordCount = sample.wordCount(wordId);
		            int topicWordCount = sample.topicWordCount(topic,wordId);
		            double topicWordProb = sample.topicWordProb(topic,wordId);
		            double z = binomialZ(topicWordCount, topicCount, wordCount, sample.numTokens());
		            
		            System.out.println();
		            System.out.printf("%6d  %15s  %7d   %4.3f  %4.3f", wordId, word, topicWordCount, topicWordProb, z);
		            
		        }
    		}
    		
    		//order word according to their z probability
    		if(z){
    			int topicCount = sample.topicCount(topic);
    			
	    		ObjectToDoubleMap<Integer> counter = new ObjectToDoubleMap<Integer>();
	    		for (int word = 0; word < sample.numWords(); word++){
	    			int topicWordCount = sample.topicWordCount(topic,word);
	    			int wordCount = sample.wordCount(word);
	    			double z = binomialZ(topicWordCount, topicCount, wordCount, sample.numTokens());
	    			counter.set(word,z);
	    		}
	    		
	    		List<Integer> listWords = counter.keysOrderedByValueList();
	    		
	    		
		        System.out.println("\nTOPIC " + topic  + "  (total count=" + topicCount + ")");
		        System.out.println("SYMBOL             WORD    COUNT   PROB     Z");
		        System.out.println("---------------------------------------------");
		            	        
		        for (int rank = 0; rank < listWords.size(); rank++) {
			    	
		        	int wordId = listWords.get(rank);
		            String word = mSymbolTable.idToSymbol(wordId);
		            int wordCount = sample.wordCount(wordId);
		            int topicWordCount = sample.topicWordCount(topic,wordId);
		            double topicWordProb = sample.topicWordProb(topic,wordId);
		            double z = binomialZ(topicWordCount, topicCount, wordCount, sample.numTokens());
		            
		            System.out.println();
		            System.out.printf("%6d  %15s  %7d   %4.3f  %4.3f", wordId, word, topicWordCount, topicWordProb, z);
		            
		        }
    		}
    		
    		//order word according to their z probability
    		if(p_z){
    			int topicCount = sample.topicCount(topic);
    			
	    		ObjectToDoubleMap<Integer> counter = new ObjectToDoubleMap<Integer>();
	    		for (int word = 0; word < sample.numWords(); word++){
	    			int topicWordCount = sample.topicWordCount(topic,word);
	    			int wordCount = sample.wordCount(word);
	    			double z = binomialZ(topicWordCount, topicCount, wordCount, sample.numTokens());
	    			
	    			double p = (double)sample.topicWordProb(topic, word);
	    			
	    			counter.set(word,z*p);
	    		}
	    		
	    		List<Integer> listWords = counter.keysOrderedByValueList();
	    		
	    		
		        System.out.println("\nTOPIC " + topic  + "  (total count=" + topicCount + ")");
		        System.out.println("SYMBOL             WORD    COUNT   PROB     Z");
		        System.out.println("---------------------------------------------");
		            	        
		        for (int rank = 0; rank < listWords.size(); rank++) {
			    	
		        	int wordId = listWords.get(rank);
		            String word = mSymbolTable.idToSymbol(wordId);
		            int wordCount = sample.wordCount(wordId);
		            int topicWordCount = sample.topicWordCount(topic,wordId);
		            double topicWordProb = sample.topicWordProb(topic,wordId);
		            double z = binomialZ(topicWordCount, topicCount, wordCount, sample.numTokens());
		            
		            System.out.println();
		            System.out.printf("%6d  %15s  %7d   %4.3f  %4.3f", wordId, word, topicWordCount, topicWordProb, z);
		            
		        }
    		}
    	}
    }
    
    public void printKml(LatentDirichletAllocation.GibbsSample sample, int topic, PrintWriter kml, Map<String,Double> labelKeyProb, Map<String,Integer> labelKeyHour) {
    	try{
    		
    		Kml kmlutil = new Kml();
    		KMLCirclePlacemark circle = new KMLCirclePlacemark();
	        
	        int mostSimilar = similarity(sample.lda(), topic);
				
	        
	        kml.println("<Folder>");
			kml.println("<name> topic "+topic+" most equal "+mostSimilar+"</name>");
	        
	

			
			for(String i: labelKeyProb.keySet()){
            	LatLonPoint center = par.mapPoints.get(i);
            	kml.println(circle.draw(center.getLongitude(), center.getLatitude(), labelKeyProb.get(i), 23, par.mapTimeIntervals.get(labelKeyHour.get(i)) , "aa0000ff", par.mapTimeIntervals.get(labelKeyHour.get(i))));
            }
			
			double [] dayweekprob = computeWeekProb(sample, topic);	            
	            
	        KMLBarChart barchart = new KMLBarChart();
	        String axis = "0:|Sun|Mon|Tue|Wed|Thu|Fri|Sat|";
	        String chart = barchart.printBarChart(dayweekprob, axis);
	        kmlutil.printScreenOverlay(kml, chart, "DOWdistribution");            
	        kml.println("</Folder>");
	        
    	}catch(Exception e){
    		e.printStackTrace();
    	}
    }
    
    public void printCsv(LatentDirichletAllocation.GibbsSample sample, int topic, Map<String,Integer> labelKeyHour) {
    	
    	ObjectToCounterMap<Integer> counter = new ObjectToCounterMap<Integer>();
    	for (int doc = 0; doc < sample.numDocuments(); doc++) 
    		counter.set(doc, sample.documentTopicCount(doc,topic));
    		
    	List<Integer> topDocs = counter.keysOrderedByCountList();

        for (int rank = 0; rank < topDocs.size(); rank++) {
        	int doc = topDocs.get(rank);
            double docTopicProb = sample.documentTopicProb(doc,topic);
            for(String h: labelKeyHour.keySet())
            	csv.println(topic+";"+par.int2filename.get(doc)+";"+docTopicProb+";"+labelKeyHour.get(h)+";"+par.mapPoints.get(h).getLatitude()+";"+par.mapPoints.get(h).getLongitude());
        }	
    	
    }

    public double [] computeWeekProb(LatentDirichletAllocation.GibbsSample sample, int topic){
        //compute days of week probabilities and best days for the given topic
        double[] dayweekprob = new double[7];
        
        DateTime dt = new DateTime();
        
        for (int doc = 0; doc < sample.numDocuments(); doc++) {
        	double docTopicProb = sample.documentTopicProb(doc, topic);
        	String dayofweek = par.int2filename.get(doc).split(",")[0];
        	
        	dayweekprob[(Integer)dt.DayOfWeekTranslator.getKey(dayofweek)-1] += docTopicProb*100;
        	
        }
        
        for(int i=0;i<dayweekprob.length;i++)
        	dayweekprob[i] = (dayweekprob[i]/((double)sample.numDocuments()/7));
        
        
        return dayweekprob;
    
    }  

    
    
	public int similarity(LatentDirichletAllocation lda, int topic){
		List<TopicSim> pairs = new ArrayList<TopicSim>();
		
		double [] pi = lda.wordProbabilities(topic);
		
		for(int i=0; i<lda.numTopics();i++){
			double [] pj = lda.wordProbabilities(i);
			double divergence = Statistics.symmetrizedKlDivergence(pi,pj);
			TopicSim ts = new TopicSim(i,topic,divergence);
			if(i != topic)
				pairs.add(ts);				
		}		
		
		Collections.sort(pairs, ScoredObject.comparator());
		
		return pairs.get(0).mI;
		
	}
    
    public double binomialZ(double wordCountInDoc, double wordsInDoc, double wordCountinCorpus, double wordsInCorpus) {
    	double pCorpus = wordCountinCorpus / wordsInCorpus;
		double var = wordsInCorpus * pCorpus * (1 - pCorpus);
		double dev = Math.sqrt(var);
		double expected = wordsInDoc * pCorpus;
		double z = (wordCountInDoc - expected) / dev;
		return z;
    }
}